
# Requirements
##Installs
def check_dependencies_already_installed() -> bool:
    try:
        import gdown
        return True
    except ImportError:
        return False


def install_dependencies():
    !pip install 'tensorflow_addons' 'tensorflow-determinism' 'gdown'


dependencies_already_installed = check_dependencies_already_installed()
print(f'dependencies_already_installed: {dependencies_already_installed}')
if not dependencies_already_installed:
    install_dependencies()

##Imports
import os
import random
import json
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.python.data.ops.dataset_ops import Dataset
from tensorflow.python.distribute.tpu_strategy import TPUStrategy
import tensorflow_addons as tfa
from skimage.metrics import structural_similarity
from scipy.stats import wasserstein_distance
import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm
from PIL import Image
from pathlib import Path
from datetime import datetime
from zipfile import ZipFile


# Environment setup
##Set tensorflow deterministic mode
def set_tf_deterministic_mode():
    os.environ['TF_DETERMINISTIC_OPS'] = '1'
    os.environ['TF_CUDNN_DETERMINISTIC'] = '1'


##Set random seed
def set_training_random_seed(seed: int):
    random.seed(seed)
    tf.random.set_seed(seed)
    np.random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    print(f'set_training_random_seed() - seed value: {seed}')


##Connect to strongest available device
set_tf_deterministic_mode()


def choose_strongest_available_device_strategy():
    try:
        tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
        tf.config.experimental_connect_to_cluster(tpu)
        tf.tpu.experimental.initialize_tpu_system(tpu)
        selected_strategy = TPUStrategy(tpu)
    except:
        selected_strategy = tf.distribute.get_strategy()

    print(f"choose_strongest_available_device_strategy() - selected strategy type: {type(selected_strategy).__name__}")

    # todo itay - delete this section so it won't mess up in google colab
    gpu_is_available = any(tf.config.list_physical_devices('GPU'))
    if gpu_is_available:
        !nvidia-smi

    return selected_strategy


DEVICE_STRATEGY = choose_strongest_available_device_strategy()

# todo itay - delete
assert tf.__version__ == '2.6.4'

##Download competition dataset
LOCAL_DATASET_FOLDER_PATH = Path('./train_data/gan-getting-started/')


def download_competition_dataset_if_not_present():
    dataset_already_downloaded = LOCAL_DATASET_FOLDER_PATH.exists()
    print(f"dataset_already_downloaded: {dataset_already_downloaded}")
    if not dataset_already_downloaded:
        # note - this is the untouched competition dataset. just uploaded it to the drive so it'll be
        # available via colab as well.
        !gdown '1ZwcoO11NKhYsbuM7hzdSzKjGOnOx6X94'
        !mkdir -p {LOCAL_DATASET_FOLDER_PATH}
        !unzip -o -q ./gan-getting-started.zip -d {LOCAL_DATASET_FOLDER_PATH}


download_competition_dataset_if_not_present()


def _choose_30_images(
        original_ordered_monet_dataset: Dataset, method: str,
        experiment_random_seed: int, use_preprocessed_cache: bool
) -> Dataset:
    set_training_random_seed(seed=42)

    try:
        original_ordered_monet_images = list(original_ordered_monet_dataset)
        if use_preprocessed_cache:
            preprocessed_indices_cache = {
                'random_selection': [203, 266, 152, 9, 233, 226, 196, 109, 5, 175, 237, 57, 218, 45, 182, 221, 289, 211, 148, 165, 78, 113, 249, 250, 104, 42, 281, 295, 157, 238],
                'farthest_images_by_pixel_distance': [57, 299, 113, 74, 160, 193, 139, 283, 93, 56, 90, 196, 108, 34, 45, 40, 249, 87, 240, 106, 218, 208, 2, 289, 16, 11, 155, 52, 292, 272],
                'closest_images_by_pixel_distance': [57, 119, 185, 203, 70, 216, 238, 61, 273, 169, 81, 222, 68, 137, 58, 67, 295, 164, 50, 217, 179, 190, 125, 290, 44, 251, 274, 195, 7, 165],
                'farthest_images_by_structural_distance': [57, 134, 16, 152, 133, 116, 94, 147, 177, 161, 40, 9, 218, 110, 70, 101, 75, 54, 243, 100, 98, 237, 77, 115, 106, 119, 45, 261, 241, 90],
                'closest_images_by_structural_distance': [57, 254, 238, 160, 86, 295, 187, 51, 41, 81, 273, 120, 230, 148, 25, 13, 264, 114, 236, 155, 58, 59, 151, 91, 107, 207, 289, 203, 176, 125],
                'farthest_images_by_earth_movers_distance': [57, 168, 144, 249, 145, 117, 288, 109, 48, 207, 252, 251, 186, 278, 296, 39, 263, 234, 167, 165, 128, 41, 290, 176, 33, 107, 202, 201, 282, 127],
                'closest_images_by_earth_movers_distance': [57, 196, 13, 131, 148, 97, 187, 74, 11, 162, 147, 240, 55, 70, 102, 76, 213, 139, 256, 253, 247, 17, 227, 108, 86, 133, 19, 77, 42, 1],
            }
            assert method in preprocessed_indices_cache, f"unknown method - '{method}'"
            chosen_30_images_indices = preprocessed_indices_cache[method]
        elif method == 'random_selection':
            chosen_30_images_indices = _pick_random_images(
                original_ordered_monet_images, images_count=30
            )
        elif method == 'farthest_images_by_pixel_distance':
            chosen_30_images_indices = _pick_images_farthest_from_each_other(
                original_ordered_monet_images,
                distance_func=_images_pixel_distance,
                images_count=30
            )
        elif method == 'closest_images_by_pixel_distance':
            chosen_30_images_indices = _pick_images_farthest_from_each_other(
                original_ordered_monet_images,
                distance_func=_images_pixel_distance,
                images_count=30,
                reverse_distance=True
            )
        elif method == 'farthest_images_by_structural_distance':
            chosen_30_images_indices = _pick_images_farthest_from_each_other(
                original_ordered_monet_images,
                distance_func=_images_structural_distance,
                images_count=30
            )
        elif method == 'closest_images_by_structural_distance':
            chosen_30_images_indices = _pick_images_farthest_from_each_other(
                original_ordered_monet_images,
                distance_func=_images_structural_distance,
                images_count=30,
                reverse_distance=True
            )
        elif method == 'farthest_images_by_earth_movers_distance':
            chosen_30_images_indices = _pick_images_farthest_from_each_other(
                original_ordered_monet_images,
                distance_func=_earth_movers_distance,
                images_count=30
            )
        elif method == 'closest_images_by_earth_movers_distance':
            chosen_30_images_indices = _pick_images_farthest_from_each_other(
                original_ordered_monet_images,
                distance_func=_earth_movers_distance,
                images_count=30,
                reverse_distance=True
            )
        else:
            raise NotImplementedError(f"unknown method - '{method}'")

        chosen_30_images_dataset = Dataset.from_tensor_slices([
            original_ordered_monet_images[image_idx]
            for image_idx in chosen_30_images_indices
        ])
        _plot_chosen_30_images(chosen_30_images_dataset)
    finally:
        set_training_random_seed(experiment_random_seed)
    return chosen_30_images_dataset


def _plot_chosen_30_images(chosen_30_images_dataset):
    try:
        images_shape = list(chosen_30_images_dataset)[0].shape
        print(f'*** Selected 30 train monet photos (shape: {images_shape}) ***')
        _, ax = plt.subplots(30, 1, figsize=(50, 50))
        for i, img in enumerate(chosen_30_images_dataset):
            img = (img * 127.5 + 127.5).numpy()[0].astype(np.uint8)

            ax[i].imshow(img)
        plt.show()
    finally:
        plt.close()


# Load competition dataset
##Load full dataset
def find_competition_dataset_files(local_dataset_folder_path: Path):
    if isinstance(DEVICE_STRATEGY, TPUStrategy):
        from kaggle_datasets import KaggleDatasets
        dataset_folder_path = Path(KaggleDatasets().get_gcs_path())
    else:
        dataset_folder_path = local_dataset_folder_path

    monet_dataset_files = tf.io.gfile.glob(str(dataset_folder_path / 'monet_tfrec/*.tfrec'))
    photo_dataset_files = tf.io.gfile.glob(str(dataset_folder_path / 'photo_tfrec/*.tfrec'))
    assert any(monet_dataset_files)
    assert any(photo_dataset_files)
    print(f"found {len(monet_dataset_files)} monet and {len(photo_dataset_files)} photo tfrec files.")

    return monet_dataset_files, photo_dataset_files


def _prepare_image_tensor_for_training(image):
    image = (tf.cast(image, tf.float32) / 127.5) - 1
    image = tf.reshape(image, [256, 256, 3])
    image = tf.image.resize(image, (320, 320), method='bilinear')
    return image


def load_tf_records_dataset(tf_record_files) -> Dataset:
    def _read_and_normalize_tfrecord(record):
        tfrecord_format = {
            "image_name": tf.io.FixedLenFeature([], tf.string),
            "image": tf.io.FixedLenFeature([], tf.string),
            "target": tf.io.FixedLenFeature([], tf.string)
        }
        record = tf.io.parse_single_example(record, tfrecord_format)
        image = record['image']
        image = tf.image.decode_jpeg(image, channels=3)
        image = _prepare_image_tensor_for_training(image)
        return image

    sorted_tf_record_files = sorted(tf_record_files)
    dataset = tf.data.TFRecordDataset(sorted_tf_record_files)
    dataset = dataset.map(_read_and_normalize_tfrecord, num_parallel_calls=tf.data.experimental.AUTOTUNE)
    return dataset


##Pick 30 train monet images strategies
def _pick_images_farthest_from_each_other(
        original_ordered_monet_images: list, distance_func, images_count: int, reverse_distance: bool = False
) -> list:
    comparison_images_resize_shape = (100, 100)

    def _pre_comparison_transformation_func(image_tensor):
        image_array = image_tensor.numpy()
        denormalized_image = (image_array * 127.5 + 127.5)
        resized_image = tf.image.resize(denormalized_image, comparison_images_resize_shape).numpy()
        resized_image = resized_image.astype(np.uint8)
        resized_image = resized_image[0]
        return resized_image

    final_distance_func = distance_func
    if reverse_distance:
        final_distance_func = lambda image1, image2: -distance_func(image1, image2)

    chosen_30_images_indices = _incremental_farthest_search(
        original_ordered_monet_images,
        k=images_count,
        distance_func=final_distance_func,
        pre_comparison_transformation_func=_pre_comparison_transformation_func
    )
    return chosen_30_images_indices


def _pick_random_images(original_ordered_monet_images: list, images_count: int) -> list:
    chosen_30_images_indices = list(np.random.choice(
        list(range(len(original_ordered_monet_images))), size=images_count, replace=False
    ))
    return chosen_30_images_indices


def _incremental_farthest_search(
        ordered_image_tensors_list, k: int, distance_func, pre_comparison_transformation_func
):
    remaining_images = [
        dict(
            orig_image_index=orig_image_index,
            img_tensor=img_tensor,
            img_comparison_array=pre_comparison_transformation_func(img_tensor)
        )
        for orig_image_index, img_tensor in enumerate(ordered_image_tensors_list)
    ]

    chosen_30_images_indices = [remaining_images.pop(random.randint(0, len(remaining_images) - 1))]
    for _ in tqdm(list(range(k - 1)), desc='incremental_farthest_search() main loop'):
        distances = [
            distance_func(
                i['img_comparison_array'],
                chosen_30_images_indices[0]['img_comparison_array']
            )
            for i in remaining_images
        ]
        for i, p in enumerate(remaining_images):
            for j, s in enumerate(chosen_30_images_indices):
                distances[i] = min(distances[i], distance_func(
                    p['img_comparison_array'], s['img_comparison_array']
                ))
        chosen_30_images_indices.append(remaining_images.pop(distances.index(max(distances))))
    chosen_30_images_indices = [i['orig_image_index'] for i in chosen_30_images_indices]
    return chosen_30_images_indices


def _images_pixel_distance(image1: np.array, image2: np.array) -> float:
    distance = np.sum((image1.flatten() - image2.flatten()) ** 2)
    return distance


def _images_structural_distance(image1: np.array, image2: np.array) -> float:
    similarity_index, *_ = structural_similarity(image1.flatten(), image2.flatten(), full=True)
    return -similarity_index


def _earth_movers_distance(image1: np.array, image2: np.array) -> float:
    image1_hist = _calc_image_greyscale_histogram(image1)
    image2_hist = _calc_image_greyscale_histogram(image2)
    distance = wasserstein_distance(image1_hist, image2_hist)
    return distance


def _calc_image_greyscale_histogram(image: np.array) -> np.array:
    greyscale_image = np.array(Image.fromarray(image).convert('L'))
    h, w = greyscale_image.shape
    hist = [0.0] * 256
    for i in range(h):
        for j in range(w):
            hist[greyscale_image[i, j]] += 1
    return np.array(hist) / (h * w)


def down_sample(filters, size, strides=2, padding='same'):
    initializer = tf.random_normal_initializer(0., 0.02)

    network = keras.Sequential()
    network.add(layers.Conv2D(filters, size, strides=strides, padding=padding,
                              kernel_initializer=initializer, use_bias=False))

    network.add(layers.LeakyReLU())

    return network


def up_sample(filters, size, strides=2, padding='same', apply_dropout=False):
    network = keras.Sequential()
    network.add(layers.Conv2DTranspose(
        filters, size, strides=strides, padding=padding, use_bias=False,
        kernel_initializer=tf.random_normal_initializer(0., 0.02)
    ))
    if apply_dropout:
        network.add(layers.Dropout(0.5))
    network.add(layers.ReLU())
    return network


def build_generator_model():
    inputs = layers.Input(shape=[320, 320, 3])

    # bs = batch size
    down_stack = [
        down_sample(64, 4),  # (bs, 160, 160, 64)
        down_sample(128, 4),  # (bs, 80, 80, 128)
        down_sample(256, 4),  # (bs, 40, 40, 256)
        down_sample(512, 4),  # (bs, 20, 20, 512)
        down_sample(512, 4),  # (bs, 10, 10, 512)
        down_sample(512, 4),  # (bs, 5, 5, 512)
        down_sample(512, 4, strides=1, padding='valid'),  # (bs, 2, 2, 512)
        down_sample(512, 4),  # (bs, 1, 1, 512)
    ]

    up_stack = [
        up_sample(512, 4, apply_dropout=True),  # (bs, 2, 2, 1024)
        up_sample(512, 4, strides=1, padding='valid', apply_dropout=True),  # (bs, 5, 5, 1024)
        up_sample(512, 4, apply_dropout=True),  # (bs, 10, 10, 1024)
        up_sample(512, 4),  # (bs, 20, 20, 1024)
        up_sample(256, 4),  # (bs, 40, 40, 512)
        up_sample(128, 4),  # (bs, 80, 80, 256)
        up_sample(64, 4),  # (bs, 160, 160, 128)
    ]

    initializer = tf.random_normal_initializer(0., 0.02)
    last = layers.Conv2DTranspose(3, 4,
                                  strides=2,
                                  padding='same',
                                  kernel_initializer=initializer,
                                  activation='tanh')  # (bs, 320, 320, 3)

    x = inputs

    # Downsampling through the model
    skips = []
    for down in down_stack:
        x = down(x)
        skips.append(x)

    skips = reversed(skips[:-1])

    # Upsampling and establishing the skip connections
    for up, skip in zip(up_stack, skips):
        x = up(x)
        x = layers.Concatenate()([x, skip])

    x = last(x)

    generator_model = keras.Model(inputs=inputs, outputs=x)
    return generator_model


"""# Build the discriminator

The discriminator takes in the input image and classifies it as real or fake (generated). Instead of outputing a single node, the discriminator outputs a smaller 2D image with higher pixel values indicating a real classification and lower values indicating a fake classification.
"""


def build_discriminator_model():
    initializer = tf.random_normal_initializer(0., 0.02)
    inp = layers.Input(shape=[320, 320, 3], name='input_image')

    x = inp
    down1 = down_sample(64, 5)(x)  # (bs, 160, 160, 64)
    down2 = down_sample(128, 4)(down1)  # (bs, 80, 80, 128)
    down3 = down_sample(256, 3)(down2)  # (bs, 40, 40, 256)
    down4 = down_sample(256, 2)(down3)  # (bs, 20, 20, 256)

    zero_pad1 = layers.ZeroPadding2D()(down4)  # (bs, 34, 34, 256)
    conv = layers.Conv2D(512, 4, strides=1,
                         kernel_initializer=initializer,
                         use_bias=False)(zero_pad1)  # (bs, 31, 31, 512)

    leaky_relu = layers.LeakyReLU()(conv)
    zero_pad2 = layers.ZeroPadding2D()(leaky_relu)  # (bs, 33, 33, 512)
    last = layers.Conv2D(1, 4, strides=1,
                         kernel_initializer=initializer)(zero_pad2)  # (bs, 30, 30, 1)

    discriminator_model = tf.keras.Model(inputs=inp, outputs=last)
    return discriminator_model


class CycleGan(keras.Model):
    def __init__(
            self,
            monet_generator,
            photo_generator,
            monet_discriminator,
            photo_discriminator,
            lambda_cycle=10,
    ):
        super(CycleGan, self).__init__()
        self.m_gen = monet_generator
        self.p_gen = photo_generator
        self.m_disc = monet_discriminator
        self.p_disc = photo_discriminator
        self.lambda_cycle = lambda_cycle

    def compile(
            self,
            m_gen_optimizer,
            p_gen_optimizer,
            m_disc_optimizer,
            p_disc_optimizer,
            gen_loss_fn,
            disc_loss_fn,
            cycle_loss_fn,
            identity_loss_fn
    ):
        super(CycleGan, self).compile()
        self.m_gen_optimizer = m_gen_optimizer
        self.p_gen_optimizer = p_gen_optimizer
        self.m_disc_optimizer = m_disc_optimizer
        self.p_disc_optimizer = p_disc_optimizer
        self.gen_loss_fn = gen_loss_fn
        self.disc_loss_fn = disc_loss_fn
        self.cycle_loss_fn = cycle_loss_fn
        self.identity_loss_fn = identity_loss_fn
        self._my_hack_history = []

    def train_step(self, batch_data):
        real_monet, real_photo = batch_data

        with tf.GradientTape(persistent=True) as tape:
            # photo to monet back to photo
            fake_monet = self.m_gen(real_photo, training=True)
            cycled_photo = self.p_gen(fake_monet, training=True)

            # monet to photo back to monet
            fake_photo = self.p_gen(real_monet, training=True)
            cycled_monet = self.m_gen(fake_photo, training=True)

            # generating itself
            same_monet = self.m_gen(real_monet, training=True)
            same_photo = self.p_gen(real_photo, training=True)

            # discriminator used to check, inputing real images
            disc_real_monet = self.m_disc(real_monet, training=True)
            disc_real_photo = self.p_disc(real_photo, training=True)

            # discriminator used to check, inputing fake images
            disc_fake_monet = self.m_disc(fake_monet, training=True)
            disc_fake_photo = self.p_disc(fake_photo, training=True)

            # evaluates generator loss
            monet_gen_loss = self.gen_loss_fn(disc_fake_monet)
            photo_gen_loss = self.gen_loss_fn(disc_fake_photo)

            # evaluates total cycle consistency loss
            total_cycle_loss = self.cycle_loss_fn(real_monet, cycled_monet, self.lambda_cycle) + self.cycle_loss_fn(
                real_photo, cycled_photo, self.lambda_cycle)

            # evaluates total generator loss
            total_monet_gen_loss = monet_gen_loss + total_cycle_loss + self.identity_loss_fn(real_monet, same_monet,
                                                                                             self.lambda_cycle)
            total_photo_gen_loss = photo_gen_loss + total_cycle_loss + self.identity_loss_fn(real_photo, same_photo,
                                                                                             self.lambda_cycle)

            # evaluates discriminator loss
            monet_disc_loss = self.disc_loss_fn(disc_real_monet, disc_fake_monet)
            photo_disc_loss = self.disc_loss_fn(disc_real_photo, disc_fake_photo)

        # Calculate the gradients for generator and discriminator
        monet_generator_gradients = tape.gradient(total_monet_gen_loss,
                                                  self.m_gen.trainable_variables)
        photo_generator_gradients = tape.gradient(total_photo_gen_loss,
                                                  self.p_gen.trainable_variables)

        monet_discriminator_gradients = tape.gradient(monet_disc_loss,
                                                      self.m_disc.trainable_variables)
        photo_discriminator_gradients = tape.gradient(photo_disc_loss,
                                                      self.p_disc.trainable_variables)

        # Apply the gradients to the optimizer
        self.m_gen_optimizer.apply_gradients(zip(monet_generator_gradients,
                                                 self.m_gen.trainable_variables))

        self.p_gen_optimizer.apply_gradients(zip(photo_generator_gradients,
                                                 self.p_gen.trainable_variables))

        self.m_disc_optimizer.apply_gradients(zip(monet_discriminator_gradients,
                                                  self.m_disc.trainable_variables))

        self.p_disc_optimizer.apply_gradients(zip(photo_discriminator_gradients,
                                                  self.p_disc.trainable_variables))

        rett = {
            "monet_gen_loss": total_monet_gen_loss,
            "photo_gen_loss": total_photo_gen_loss,
            "monet_disc_loss": monet_disc_loss,
            "photo_disc_loss": photo_disc_loss
        }
        self._my_hack_history.append(rett)
        return rett


#Graph plotting utils
def plot_cycle_gan_train_losses(train_history):
    epoch_level_train_history = {
        loss_name: {
            epoch: epoch_losses.flatten().mean()
            for epoch, epoch_losses in enumerate(epochs_losses)
        }
        for loss_name, epochs_losses in train_history.history.items()
    }

    try:
        fig, (top_ax, bottom_ax) = plt.subplots(2, figsize=(12, 12))
        fig.suptitle('Loss vs Epoch')

        top_ax.set_title('Generators')
        plt_lines = []
        for loss_name in ('monet_gen_loss', 'photo_gen_loss'):
            loss_epoch_values = epoch_level_train_history[loss_name]
            line, = top_ax.plot(loss_epoch_values.keys(), loss_epoch_values.values(), label=loss_name)
            plt_lines.append(line)
        top_ax.legend(handles=plt_lines)

        bottom_ax.set_title('Discriminators')
        plt_lines = []
        for loss_name in ('monet_disc_loss', 'photo_disc_loss'):
            loss_epoch_values = epoch_level_train_history[loss_name]
            line, = bottom_ax.plot(loss_epoch_values.keys(), loss_epoch_values.values(), label=loss_name)
            plt_lines.append(line)
        bottom_ax.legend(handles=plt_lines)

        plt.tight_layout()
        plt.show()
    finally:
        plt.close()

    final_losses = {
        loss_name: str(loss_epoch_values[max(loss_epoch_values.keys())])
        for loss_name, loss_epoch_values in epoch_level_train_history.items()
    }
    print(f"*** trained cycle gan final losses ***\n{json.dumps(final_losses, indent=4)}")


"""# Define loss functions

The discriminator loss function below compares real images to a matrix of 1s and fake images to a matrix of 0s. The perfect discriminator will output all 1s for real images and all 0s for fake images. The discriminator loss outputs the average of the real and generated loss.
"""

"""# Create submission file"""
def create_predictions_for_kaggle_submission(monet_generator: tf.keras.Model, photo_dataset: Dataset):
    wip_images_folder = Path(f"/tmp/wip_{datetime.now().strftime('%y_%m_%d__%H_%M_%S')}/")
    wip_images_folder.mkdir(parents=True, exist_ok=True)
    output_zip_path = Path('/kaggle/working/images.zip')
    if output_zip_path.exists():
        output_zip_path.unlink()

    with ZipFile(output_zip_path, 'w') as output_images_zip, \
            tqdm(total=7_038, desc='generating prediction images for kaggle submission') as pbar:
        for i, img in enumerate(photo_dataset):
            prediction = monet_generator(img, training=False)[0]
            prediction = tf.image.resize(prediction, (256, 256), method='bilinear')
            prediction = prediction.numpy()
            prediction = (prediction * 127.5 + 127.5).astype(np.uint8)
            im = Image.fromarray(prediction)
            wip_image_file_path = wip_images_folder / f'{i + 1}.jpg'
            im.save(wip_image_file_path)
            output_images_zip.write(filename=wip_image_file_path, arcname=wip_image_file_path.name)
            wip_image_file_path.unlink()
            pbar.update()


def experiment_flow(
        choose_30_images_method: str,
        train_settings: dict,
        experiment_random_seed: int,
        create_kaggle_predictions_for_submission: bool = False,
):
    set_training_random_seed(experiment_random_seed)

    monet_dataset_files, photo_dataset_files = find_competition_dataset_files(LOCAL_DATASET_FOLDER_PATH)
    original_ordered_monet_dataset = load_tf_records_dataset(monet_dataset_files).batch(1)
    photo_dataset = load_tf_records_dataset(photo_dataset_files).batch(1)

    chosen_30_monet_dataset = _choose_30_images(
        original_ordered_monet_dataset, choose_30_images_method,
        experiment_random_seed, use_preprocessed_cache=True
    )

    with DEVICE_STRATEGY.scope():
        monet_generator = build_generator_model()  # transforms photos to Monet-esque paintings
        photo_generator = build_generator_model()  # transforms Monet paintings to be more like photos

        monet_discriminator = build_discriminator_model()  # differentiates real Monet paintings and generated Monet paintings
        photo_discriminator = build_discriminator_model()  # differentiates real photos and generated photos

    with DEVICE_STRATEGY.scope():
        def discriminator_loss(real, generated):
            real_loss = tf.keras.losses.BinaryCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)(
                tf.ones_like(real), real)

            generated_loss = tf.keras.losses.BinaryCrossentropy(from_logits=True,
                                                                reduction=tf.keras.losses.Reduction.NONE)(
                tf.zeros_like(generated), generated)

            total_disc_loss = real_loss + generated_loss

            return total_disc_loss * 0.5

    """The generator wants to fool the discriminator into thinking the generated image is real. The perfect generator will have the discriminator output only 1s. Thus, it compares the generated image to a matrix of 1s to find the loss."""

    with DEVICE_STRATEGY.scope():
        def generator_loss(generated):
            return tf.keras.losses.BinaryCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)(
                tf.ones_like(generated), generated)

    """We want our original photo and the twice transformed photo to be similar to one another. Thus, we can calculate the cycle consistency loss be finding the average of their difference."""

    with DEVICE_STRATEGY.scope():
        def calc_cycle_loss(real_image, cycled_image, lambda_):
            loss1 = tf.reduce_mean(tf.abs(real_image - cycled_image))

            return lambda_ * loss1

    """The identity loss compares the image with its generator (i.e. photo with photo generator). If given a photo as input, we want it to generate the same image as the image was originally a photo. The identity loss compares the input with the output of the generator."""

    with DEVICE_STRATEGY.scope():
        def identity_loss(real_image, same_image, lambda_):
            loss = tf.reduce_mean(tf.abs(real_image - same_image))
            return lambda_ * 0.5 * loss

    """# Train the CycleGAN

    Let's compile our model. Since we used `tf.keras.Model` to build our CycleGAN, we can just ude the `fit` function to train our model.
    """

    with DEVICE_STRATEGY.scope():
        optimizer_builder = train_settings['optimizer_builder']
        monet_generator_optimizer = optimizer_builder()
        photo_generator_optimizer = optimizer_builder()

        monet_discriminator_optimizer = optimizer_builder()
        photo_discriminator_optimizer = optimizer_builder()

    with DEVICE_STRATEGY.scope():
        cycle_gan_model = CycleGan(
            monet_generator, photo_generator, monet_discriminator, photo_discriminator
        )

        cycle_gan_model.compile(
            m_gen_optimizer=monet_generator_optimizer,
            p_gen_optimizer=photo_generator_optimizer,
            m_disc_optimizer=monet_discriminator_optimizer,
            p_disc_optimizer=photo_discriminator_optimizer,
            gen_loss_fn=generator_loss,
            disc_loss_fn=discriminator_loss,
            cycle_loss_fn=calc_cycle_loss,
            identity_loss_fn=identity_loss
        )

    print(f'\n\n*** running_on_tpu - {isinstance(DEVICE_STRATEGY, TPUStrategy)} ***\n\n')

    train_history = cycle_gan_model.fit(
        tf.data.Dataset.zip((chosen_30_monet_dataset, photo_dataset)),
        epochs=train_settings['train_epochs'],
        verbose=2
    )
    plot_cycle_gan_train_losses(train_history)

    print('*** Show trained model predictions sample ***')
    _, ax = plt.subplots(5, 2, figsize=(12, 12))
    for i, img in enumerate(photo_dataset.take(5)):
        prediction = monet_generator(img, training=False)[0].numpy()
        prediction = (prediction * 127.5 + 127.5).astype(np.uint8)
        img = (img[0] * 127.5 + 127.5).numpy().astype(np.uint8)

        ax[i, 0].imshow(img)
        ax[i, 1].imshow(prediction)
        ax[i, 0].set_title("Input Photo")
        ax[i, 1].set_title("Monet-esque")
        ax[i, 0].axis("off")
        ax[i, 1].axis("off")
    plt.show()

    if create_kaggle_predictions_for_submission:
        create_predictions_for_kaggle_submission(monet_generator, photo_dataset)


choose_30_images_methods = (
    'random_selection',
    'farthest_images_by_pixel_distance',
    'closest_images_by_pixel_distance',
    'farthest_images_by_structural_distance',
    'closest_images_by_structural_distance',
    'farthest_images_by_earth_movers_distance',
    'closest_images_by_earth_movers_distance',
)

base_desc = 'choose_30_images_methods loop'
with tqdm(total=len(choose_30_images_methods), desc=base_desc) as pbar:
    for method in choose_30_images_methods:
        pbar.set_description(f"{base_desc} (curr_method = '{method}')")
        experiment_flow(
            choose_30_images_method=method,
            train_settings=dict(
                train_epochs=40,
                optimizer_builder=lambda: tf.keras.optimizers.Adam(learning_rate=0.001, decay=0.001)
            ),
            experiment_random_seed=1,
            create_kaggle_predictions_for_submission=False
        )
        pbar.update()

dependencies_already_installed: True
choose_strongest_available_device_strategy() - selected strategy type: _DefaultDistributionStrategy
Wed Feb 15 10:08:05 2023       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.82.01    Driver Version: 470.82.01    CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla P100-PCIE...  Off  | 00000000:00:04.0 Off |                    0 |
| N/A   42C    P0    36W / 250W |  15977MiB / 16280MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+
dataset_already_downloaded: True

choose_30_images_methods loop (curr_method = 'random_selection'):   0%|          | 0/7 [00:00<?, ?it/s]

set_training_random_seed() - seed value: 1
found 5 monet and 20 photo tfrec files.
set_training_random_seed() - seed value: 42

2023-02-15 10:08:06.353221: W tensorflow/core/data/root_dataset.cc:167] Optimization loop failed: Cancelled: Operation was cancelled

*** Selected 30 train monet photos (shape: (1, 320, 320, 3)) ***

set_training_random_seed() - seed value: 1


*** running_on_tpu - False ***


Epoch 1/40
30/30 - 19s - monet_gen_loss: 6.9626 - photo_gen_loss: 5.9112 - monet_disc_loss: 0.2345 - photo_disc_loss: 0.6435
Epoch 2/40
30/30 - 12s - monet_gen_loss: 6.7633 - photo_gen_loss: 4.6662 - monet_disc_loss: 0.0989 - photo_disc_loss: 0.6746
Epoch 3/40
30/30 - 12s - monet_gen_loss: 4.6350 - photo_gen_loss: 4.8057 - monet_disc_loss: 0.7258 - photo_disc_loss: 0.6197
Epoch 4/40
30/30 - 12s - monet_gen_loss: 3.5104 - photo_gen_loss: 3.1220 - monet_disc_loss: 0.7640 - photo_disc_loss: 0.7535
Epoch 5/40
30/30 - 12s - monet_gen_loss: 3.4218 - photo_gen_loss: 2.8866 - monet_disc_loss: 0.8662 - photo_disc_loss: 0.7112
Epoch 6/40
30/30 - 12s - monet_gen_loss: 3.2593 - photo_gen_loss: 2.6979 - monet_disc_loss: 0.7104 - photo_disc_loss: 0.7792
Epoch 7/40
30/30 - 12s - monet_gen_loss: 3.2395 - photo_gen_loss: 2.6063 - monet_disc_loss: 0.7459 - photo_disc_loss: 0.7625
Epoch 8/40
30/30 - 12s - monet_gen_loss: 3.1849 - photo_gen_loss: 2.5222 - monet_disc_loss: 0.8755 - photo_disc_loss: 0.8341
Epoch 9/40
30/30 - 12s - monet_gen_loss: 3.3647 - photo_gen_loss: 2.4910 - monet_disc_loss: 0.8527 - photo_disc_loss: 0.8187
Epoch 10/40
30/30 - 12s - monet_gen_loss: 3.3144 - photo_gen_loss: 2.7887 - monet_disc_loss: 0.8126 - photo_disc_loss: 0.6342
Epoch 11/40
30/30 - 12s - monet_gen_loss: 3.0709 - photo_gen_loss: 2.5117 - monet_disc_loss: 0.8032 - photo_disc_loss: 0.7544
Epoch 12/40
30/30 - 12s - monet_gen_loss: 3.3065 - photo_gen_loss: 2.7014 - monet_disc_loss: 0.9114 - photo_disc_loss: 0.7117
Epoch 13/40
30/30 - 12s - monet_gen_loss: 3.2505 - photo_gen_loss: 2.4335 - monet_disc_loss: 0.8669 - photo_disc_loss: 0.8110
Epoch 14/40
30/30 - 12s - monet_gen_loss: 3.0201 - photo_gen_loss: 2.3520 - monet_disc_loss: 0.8344 - photo_disc_loss: 0.8064
Epoch 15/40
30/30 - 12s - monet_gen_loss: 3.1055 - photo_gen_loss: 2.3546 - monet_disc_loss: 0.9768 - photo_disc_loss: 0.8247
Epoch 16/40
30/30 - 12s - monet_gen_loss: 2.8504 - photo_gen_loss: 2.3155 - monet_disc_loss: 0.8964 - photo_disc_loss: 0.7060
Epoch 17/40
30/30 - 12s - monet_gen_loss: 3.0314 - photo_gen_loss: 2.3105 - monet_disc_loss: 0.9045 - photo_disc_loss: 0.8191
Epoch 18/40
30/30 - 12s - monet_gen_loss: 3.1208 - photo_gen_loss: 2.4715 - monet_disc_loss: 0.9147 - photo_disc_loss: 0.7057
Epoch 19/40
30/30 - 12s - monet_gen_loss: 2.8780 - photo_gen_loss: 2.1129 - monet_disc_loss: 0.9312 - photo_disc_loss: 0.9064
Epoch 20/40
30/30 - 12s - monet_gen_loss: 2.9279 - photo_gen_loss: 2.4213 - monet_disc_loss: 0.9045 - photo_disc_loss: 0.6794
Epoch 21/40
30/30 - 12s - monet_gen_loss: 2.9812 - photo_gen_loss: 2.2528 - monet_disc_loss: 0.9079 - photo_disc_loss: 0.8212
Epoch 22/40
30/30 - 12s - monet_gen_loss: 2.8172 - photo_gen_loss: 2.1002 - monet_disc_loss: 0.8932 - photo_disc_loss: 0.8550
Epoch 23/40
30/30 - 12s - monet_gen_loss: 2.7932 - photo_gen_loss: 2.2246 - monet_disc_loss: 0.9597 - photo_disc_loss: 0.7914
Epoch 24/40
30/30 - 12s - monet_gen_loss: 2.8371 - photo_gen_loss: 2.2055 - monet_disc_loss: 0.7947 - photo_disc_loss: 0.8277
Epoch 25/40
30/30 - 12s - monet_gen_loss: 3.0438 - photo_gen_loss: 2.4541 - monet_disc_loss: 0.9631 - photo_disc_loss: 0.7323
Epoch 26/40
30/30 - 12s - monet_gen_loss: 2.8756 - photo_gen_loss: 2.1502 - monet_disc_loss: 0.8776 - photo_disc_loss: 0.8439
Epoch 27/40
30/30 - 12s - monet_gen_loss: 2.9542 - photo_gen_loss: 2.4061 - monet_disc_loss: 0.9210 - photo_disc_loss: 0.6947
Epoch 28/40
30/30 - 12s - monet_gen_loss: 2.8047 - photo_gen_loss: 2.0609 - monet_disc_loss: 0.8612 - photo_disc_loss: 0.8932
Epoch 29/40
30/30 - 12s - monet_gen_loss: 2.8905 - photo_gen_loss: 2.3097 - monet_disc_loss: 0.9428 - photo_disc_loss: 0.7569
Epoch 30/40
30/30 - 12s - monet_gen_loss: 2.7459 - photo_gen_loss: 2.1800 - monet_disc_loss: 0.8869 - photo_disc_loss: 0.8299
Epoch 31/40
30/30 - 12s - monet_gen_loss: 2.8500 - photo_gen_loss: 2.3103 - monet_disc_loss: 0.8642 - photo_disc_loss: 0.7043
Epoch 32/40
30/30 - 12s - monet_gen_loss: 2.8857 - photo_gen_loss: 2.1102 - monet_disc_loss: 0.9202 - photo_disc_loss: 0.8771
Epoch 33/40
30/30 - 12s - monet_gen_loss: 2.7054 - photo_gen_loss: 2.2203 - monet_disc_loss: 0.8380 - photo_disc_loss: 0.7105
Epoch 34/40
30/30 - 12s - monet_gen_loss: 2.7378 - photo_gen_loss: 2.1267 - monet_disc_loss: 0.8733 - photo_disc_loss: 0.8300
Epoch 35/40
30/30 - 12s - monet_gen_loss: 2.6889 - photo_gen_loss: 2.1963 - monet_disc_loss: 0.8667 - photo_disc_loss: 0.7326
Epoch 36/40
30/30 - 12s - monet_gen_loss: 2.7621 - photo_gen_loss: 2.1806 - monet_disc_loss: 0.8704 - photo_disc_loss: 0.7980
Epoch 37/40
30/30 - 12s - monet_gen_loss: 2.7949 - photo_gen_loss: 2.1478 - monet_disc_loss: 0.9164 - photo_disc_loss: 0.7865
Epoch 38/40
30/30 - 12s - monet_gen_loss: 2.6348 - photo_gen_loss: 2.1189 - monet_disc_loss: 0.8554 - photo_disc_loss: 0.7779
Epoch 39/40
30/30 - 12s - monet_gen_loss: 2.9706 - photo_gen_loss: 2.2352 - monet_disc_loss: 1.0323 - photo_disc_loss: 0.7970
Epoch 40/40
30/30 - 12s - monet_gen_loss: 2.5987 - photo_gen_loss: 2.1137 - monet_disc_loss: 0.9298 - photo_disc_loss: 0.8303

*** trained cycle gan final losses ***
{
    "monet_gen_loss": "2.598734",
    "photo_gen_loss": "2.1136844",
    "monet_disc_loss": "0.929799",
    "photo_disc_loss": "0.83030796"
}
*** Show trained model predictions sample ***

choose_30_images_methods loop (curr_method = 'farthest_images_by_pixel_distance'):  14%|█▍        | 1/7 [13:34<1:21:29, 814.87s/it]

set_training_random_seed() - seed value: 1
found 5 monet and 20 photo tfrec files.
set_training_random_seed() - seed value: 42

2023-02-15 10:21:41.146216: W tensorflow/core/data/root_dataset.cc:167] Optimization loop failed: Cancelled: Operation was cancelled